import sys
import numpy as np
import torch
import torch.nn as nn
import ctypes

sys.path.append("../utils")

from utils import OpRunner, compare, verify_result
import time


def main():
    func = OpRunner("./build/libop.so", 4, pint_time=True)
    func.op.run.argtypes[3] = ctypes.c_char_p
    minvalue = 1
    maxvalue = 10
    shape = [1024*10*1024]
    for dtype in [np.float16, np.float32]:
        print("测试类型", dtype)
        for reduction in ["mean", "sum", "none"]:
            input_predict = np.random.uniform(minvalue, maxvalue, shape).astype(dtype)
            input_label = np.random.uniform(minvalue, maxvalue, shape).astype(dtype)
            input_dout = np.random.uniform(minvalue, maxvalue, [1]).astype(dtype)
            output = np.zeros(shape, dtype=dtype)
            if 'mean' == reduction:
                reduce_elts = 1.0
                for i in input_predict.shape:
                    reduce_elts *= i
                cof = (reduce_elts**(-1)) * 2.0
            else:
                cof = 2.0
            sub_res = input_predict - input_label
            norm_grad = sub_res * cof
            golden = norm_grad * input_dout

            func([input_predict, input_label, input_dout, reduction.encode(), output], [-1])

            print("output", output[0])
            print("golden", golden)
            verify_result(output, golden)
            print("----------------------------")


if __name__ == "__main__":
    main()
